function [spike_samples,slab_samples,prec_spike_samples,prec_slab_samples,llik_samples] = sample_spike_slab_beta_reg(y,X,varargin)
% ================================================================
%
% [spike,slab,prec_spike,prec_slab,llik] = sample_spike_slab_beta_reg(y,X,[options])
%
%  y               = n x 1 response vector (0, 1)
%  X               = n x p predictor matrix
%  burnin          = 1 x 1
%  ngibbs          = 1 x 1
%
%  spike           = p x ngibbs
%  slab            = p x ngibbs
%  prec_spike      = p x ngibbs
%  prec_slab       = p x ngibbs
%  llik            = 1 x (ngibbs + burnin)
%
%
% [options (default values)]
%
%  PRIOR_ODDS      = p x 1 prior odds for mean model
%  SLICE_WIDTH     = 0.5  slice sampling width
%  SLAB_LOWER      = -5   range of slab
%  SLAB_UPPER      = 5    range of slab
%  PREC_SLAB_LOWER = -5   range of slab in precision
%  PREC_SLAB_UPPER = 5    range of slab in precision
%  burnin          = 100  burn-in period
%  ngibbs          = 1000 number of gibbs sampling
%  verbose         = 1    for print out (set 0 to avoid output)
%
% ================================================================
%
% code: Yongjin Park, ypp@csail.mit.edu
%

    SLICE_WIDTH     = 0.5;  % default slice sampling width
    TOL             = 1e-4; % default tolerance level to avoid 1 or 0
    SLAB_LOWER      = -5;   % default range of slab
    SLAB_UPPER      = 5;    % default range of slab
    PREC_SLAB_LOWER = -5;   % default range of slab in precision
    PREC_SLAB_UPPER = 5;    % default range of slab in precision
    burnin = 100;
    ngibbs = 1000;
    verbose = 1;

    [n,p] = size(X);
    prior_odds = zeros(p,1,'single');

    while ~isempty(varargin)
        switch upper(varargin{1})
          case 'BURNIN'
            burnin          = varargin{2};
            varargin(1:2)   = [];

          case 'NGIBBS'
            ngibbs          = varargin{2};
            varargin(1:2)   = [];

          case 'SLICE_WIDTH'
            SLICE_WIDTH     = varargin{2};
            varargin(1:2)   = [];

          case 'TOL'
            TOL             = varargin{2};
            varargin(1:2)   = [];

          case 'SLAB_LOWER'
            SLAB_LOWER      = varargin{2};
            varargin(1:2)   = [];

          case 'SLAB_UPPER'
            SLAB_UPPER      = varargin{2};
            varargin(1:2)   = [];

          case 'PREC_SLAB_LOWER'
            PREC_SLAB_LOWER = varargin{2};
            varargin(1:2)   = [];

          case 'PREC_SLAB_UPPER'
            PREC_SLAB_UPPER = varargin{2};
            varargin(1:2)   = [];

          case 'PRIOR_ODDS'
            prior_odds      = varargin{2};
            varargin(1:2)   = [];

          case 'VERBOSE'
            verbose         = varargin{2};
            varargin(1:2)   = [];

          otherwise
            error(['Unexpected option: ', varargin{1}]);
        end
    end

    % ================================================================
    % add intercept
    X = [X, ones(n,1)];
    prior_odds = [prior_odds; 0];

    [~,p] = size(X);

    assert(size(y,1) == n && size(y,2) == 1, 'check if y = n x 1');
    assert(size(prior_odds,1) == p && size(prior_odds,2) == 1, 'check if prior_odds = p x 1');

    % check extreme value y = 0 or y = 1
    extreme = 1e-4;
    if any(y(:) < extreme) || any(y(:) > 1-extreme),
        y = (y - 0.5)*(1-extreme) + 0.5;
    end

    spike_samples = ones(p,ngibbs,'single');
    prec_slab_samples = zeros(p,ngibbs,'single');
    prec_spike_samples = zeros(p,ngibbs,'single');
    slab_samples = zeros(p,ngibbs,'single');
    llik_samples = zeros(1,ngibbs+burnin,'single');

    spike = ones(p,1,'single');
    slab = zeros(p,1,'single');
    prec_spike = ones(p,1,'single');
    prec_slab = zeros(p,1);
    prec = exp(X*(prec_spike.*prec_slab)); % prec ~ exp(1+delta) and delta ~ N(0,1)
    prec_slab_samples(p,ngibbs);

    for iter = 1:(burnin + ngibbs),

        % ================================================================
        % 1. Sample slab and spike by slice & gibbs sampling

        % evaluated log density at current slab
        for k = 1:p,
            fval = llik_func(y,X,spike,slab,prec);
            fval_prime = fval + log(rand);

            % (i) slice sampling of slab
            % (i-a) step-out
            left = slab;
            right = slab;
            slab_prime = slab;

            rr = rand;
            left(k) = left(k) - rr * SLICE_WIDTH;
            right(k) = right(k) + (1-rr) * SLICE_WIDTH;

            while (llik_func(y,X,spike,left,prec) > fval_prime) && left(k) > SLAB_LOWER,
                left(k) = left(k) - SLICE_WIDTH;
            end
            left(k) = max(SLAB_LOWER, left(k));

            while (llik_func(y,X,spike,right,prec) > fval_prime) && right(k) < SLAB_UPPER,
                right(k) = right(k) + SLICE_WIDTH;
            end
            right(k) = min(SLAB_UPPER, right(k));

            % (i-b) shrinkage
            while 1
                slab_prime(k) = rand() * (right(k) - left(k)) + left(k);
                fval = llik_func(y,X,spike,slab_prime,prec);

                if fval >= fval_prime - TOL,
                    break;
                else
                    if slab_prime(k) > slab(k)
                        right(k) = slab_prime(k);
                    elseif slab_prime(k) < slab(k)
                        left(k) = slab_prime(k);
                    else
                        error('error in slab sampling');
                    end
                end
            end
            slab = slab_prime;

            % (ii) sample spike by gibbs sampling
            spike(k) = 0;
            mass0 = llik_func(y,X,spike,slab_prime,prec);
            spike(k) = 1;
            mass1 = llik_func(y,X,spike,slab_prime,prec);
            pr = 1./(1+exp(- (mass1 - mass0 + prior_odds(k))));
            spike(k) = single(rand() < pr);

        end


        % ================================================================
        % 2. Sample prec_slab model
        for k = 1:p,
            fval = llik_func(y,X,spike,slab,prec);
            fval_prime = fval + log(rand);

            % (i) sample effect of displacement
            % (i-a) step out
            left = prec_slab;
            right = prec_slab;
            prec_slab_prime = prec_slab;

            rr = rand;
            left(k) = left(k) - rr * SLICE_WIDTH;
            right(k) = right(k) + (1-rr) * SLICE_WIDTH;

            prec_prime = exp(X*(prec_spike.*left));
            while (llik_func(y,X,spike,slab,prec_prime) > fval_prime) && left(k) > PREC_SLAB_LOWER,
                left(k) = left(k) - SLICE_WIDTH;
            end
            left(k) = max(PREC_SLAB_LOWER, left(k));

            prec_prime = exp(X*(prec_spike.*right));
            while (llik_func(y,X,spike,slab,prec_prime) > fval_prime) && right(k) < PREC_SLAB_UPPER,
                right(k) = right(k) + SLICE_WIDTH;
            end
            right(k) = min(PREC_SLAB_UPPER, right(k));


            % (i-b) pick a point by shrinkage
            while 1
                prec_slab_prime(k) = rand() * (right(k) - left(k)) + left(k);
                prec_prime = exp(X*(prec_spike.*prec_slab_prime));

                fval = llik_func(y,X,spike,slab,prec_prime);

                if fval >= fval_prime - TOL,
                    break;
                else
                    if prec_slab_prime(k) > prec_slab(k)
                        right(k) = prec_slab_prime(k);
                    elseif prec_slab_prime(k) < prec_slab(k)
                        left(k) = prec_slab_prime(k);
                    else
                        error('error in prec sampling');
                    end
                end
            end

            prec_slab = prec_slab_prime;
            prec = prec_prime;

            % (ii) sample spike to enforce sparsity
            prec_spike(k) = 0;
            prec_prime = exp(X*(prec_spike.*prec_slab));
            mass0 = llik_func(y,X,spike,slab,prec_prime);

            prec_spike(k) = 1;
            prec_prime = exp(X*(prec_spike.*prec_slab));
            mass1 = llik_func(y,X,spike,slab,prec_prime);

            pr = 1./(1+exp(- (mass1 - mass0)));
            prec_spike(k) = single(rand() < pr);
            prec = exp(X*(prec_spike.*prec_slab));
        end


        % keep track of samples
        if iter > burnin,
            spike_samples(:,iter - burnin) = spike(1:p);
            slab_samples(:,iter - burnin) = slab(1:p);
            prec_spike_samples(:,iter-burnin) = prec_spike;
            prec_slab_samples(:,iter-burnin) = prec_slab;
        end

        llik = llik_func(y,X,spike,slab_prime,prec);
        llik_samples(iter) = llik;
        if verbose
            fprintf('iter = %05d, llik = %.4e\r',iter,llik);
        end

        if isnan(llik) || abs(llik) == Inf,
            error(['Unusual log-likelihood = ', num2str(llik)]);
        end
    end
end

function fval = llik_func(y,X,spike,slab,prec)

    TOL = 1e-4;

    assert(all(prec + TOL > 0), 'must be positive precision');
    prec(prec < TOL) = TOL;

    mu = 1./(1+exp(- X * (spike.*slab)));
    mu(mu < TOL) = TOL;
    mu(mu > 1 - TOL) = 1 - TOL;
    a = mu.*prec;
    b = (1-mu).*prec;

    fval = sum(gammaln(prec) -gammaln(a) - gammaln(b) + (a-1).*log(y) + (b-1).*log(1-y) );
end
